import torch 
from torch import nn
import numpy as np 
from model.dataloader import get_dataset_info

word2id, id2word, vocab_size = get_dataset_info()

# 首选gpu
if torch.cuda.is_available(): 
    device = torch.device("cuda")
else:
    device = torch.device("cpu")
  
# 最大位置编码值
max_pos = 1800
# 嵌入层的大小
d_model = 768  
# 前馈网络的维度
d_ff = 2048  
# K(=Q)和V的维度
d_k = d_v = 64
# 编码器和解码器层的数量
n_layers = 6
# 多头注意力中头的数量
n_heads = 8
# 梯度裁剪的阈值
CLIP = 1

def get_attn_pad_mask(seq_q, seq_k):
    batch_size, len_q = seq_q.size()
    batch_size, len_k = seq_k.size()
    pad_attn_mask = seq_k.data.eq(0).unsqueeze(1)
    return pad_attn_mask.expand(batch_size, len_q, len_k)

def get_attn_subsequence_mask(seq):
    attn_shape = [seq.size(0), seq.size(1), seq.size(1)]

    subsequence_mask = np.triu(np.ones(attn_shape), k=1)

    subsequence_mask = torch.from_numpy(subsequence_mask).byte()

    subsequence_mask = subsequence_mask.to(device)

    return subsequence_mask 

class ScaledDotProductAttention(nn.Module):
    def __init__(self):
        super(ScaledDotProductAttention, self).__init__()  

    def forward(self, Q, K, V, attn_mask):
        scores = torch.matmul(Q, K.transpose(-1, -2)) / np.sqrt(d_k)

        scores.masked_fill_(attn_mask, -1e9)

        attn = nn.Softmax(dim=-1)(scores)

        context = torch.matmul(attn, V)

        return context, attn

class MultiHeadAttention(nn.Module):
    def __init__(self):
        super(MultiHeadAttention, self).__init__() 

        self.W_Q = nn.Linear(d_model, d_k * n_heads, bias=False)
        self.W_K = nn.Linear(d_model, d_k * n_heads, bias=False)
        self.W_V = nn.Linear(d_model, d_v * n_heads, bias=False)

        self.fc = nn.Linear(n_heads * d_v, d_model, bias=False)

        self.layernorm = nn.LayerNorm(d_model)

    def forward(self, input_Q, input_K, input_V, attn_mask):
        residual = input_Q
        batch_size = input_Q.size(0)

        Q = self.W_Q(input_Q).view(batch_size, -1, n_heads, d_k).transpose(1, 2)  
        K = self.W_K(input_K).view(batch_size, -1, n_heads, d_k).transpose(1, 2)  
        V = self.W_V(input_V).view(batch_size, -1, n_heads, d_v).transpose(1, 2)  

        attn_mask = attn_mask.unsqueeze(1).repeat(1, n_heads, 1, 1) 

        context, attn = ScaledDotProductAttention()(Q, K, V, attn_mask)

        context = context.transpose(1, 2).reshape(batch_size, -1, n_heads * d_v) 

        output = self.fc(context)  

        return self.layernorm(output + residual), attn

class PoswiseFeedForwardNet(nn.Module):
    def __init__(self):
        super(PoswiseFeedForwardNet, self).__init__() 

        self.fc = nn.Sequential(
            nn.Linear(d_model, d_ff, bias=False),
            nn.ReLU(),
            nn.Linear(d_ff, d_model, bias=False)
        )

        self.layernorm = nn.LayerNorm(d_model)

    def forward(self, inputs):
        residual = inputs
        output = self.fc(inputs)
        return self.layernorm(output + residual) 

class DecoderLayer(nn.Module):
    def __init__(self):
        super(DecoderLayer, self).__init__()  
        self.dec_self_attn = MultiHeadAttention()  
        self.dec_enc_attn = MultiHeadAttention() 
        self.pos_ffn = PoswiseFeedForwardNet()  

    def forward(self, dec_inputs, dec_self_attn_mask):
        dec_outputs, dec_self_attn = self.dec_self_attn(dec_inputs, dec_inputs, dec_inputs, dec_self_attn_mask)

        dec_outputs = self.pos_ffn(dec_outputs) 

        return dec_outputs, dec_self_attn

class Decoder(nn.Module):
    def __init__(self):
        super(Decoder, self).__init__()

        self.tgt_emb = nn.Embedding(vocab_size, d_model)

        self.pos_emb = nn.Embedding(max_pos, d_model)

        self.layers = nn.ModuleList([DecoderLayer() for _ in range(n_layers)])

    def forward(self, dec_inputs):
        seq_len = dec_inputs.size(1) 

        pos = torch.arange(seq_len, dtype=torch.long, device=device) 
        
        pos = pos.unsqueeze(0).expand_as(dec_inputs) 

        dec_outputs = self.tgt_emb(dec_inputs) + self.pos_emb(pos) 

        dec_self_attn_pad_mask = get_attn_pad_mask(dec_inputs, dec_inputs)  
        dec_self_attn_subsequence_mask = get_attn_subsequence_mask(dec_inputs)  
        dec_self_attn_mask = torch.gt((dec_self_attn_pad_mask + dec_self_attn_subsequence_mask), 0) 

        dec_self_attns = []

        for layer in self.layers:
            dec_outputs, dec_self_attn = layer(dec_outputs, dec_self_attn_mask)
            dec_self_attns.append(dec_self_attn)

        return dec_outputs, dec_self_attns

class GPT(nn.Module):
    def __init__(self):
        super(GPT, self).__init__() 

        self.decoder = Decoder() 
        self.projection = nn.Linear(d_model, vocab_size) 

    def forward(self, dec_inputs):
        dec_outputs, dec_self_attns = self.decoder(dec_inputs)  
        dec_logits = self.projection(dec_outputs)
        return dec_logits.view(-1, dec_logits.size(-1)), dec_self_attns

    def greedy_decoder(self, dec_input):
        terminal = False
        start_dec_len = len(dec_input[0])
        while not terminal:
            if len(dec_input[0]) - start_dec_len > 100:
                next_symbol = word2id['<sep>']
                dec_input = torch.cat(
                    [dec_input.detach(), torch.tensor([[next_symbol]], dtype=dec_input.dtype, device=device)], -1)
                break
            dec_outputs, _ = self.decoder(dec_input)
            projected = self.projection(dec_outputs)
            prob = projected.squeeze(0).max(dim=-1, keepdim=False)[1]
            next_word = prob.data[-1]
            next_symbol = next_word
            if next_symbol == word2id["<sep>"]:
                terminal = True
            dec_input = torch.cat(
                [dec_input.detach(), torch.tensor([[next_symbol]], dtype=dec_input.dtype, device=device)], -1)
        return dec_input

    def answer(self, sentence):
        dec_input = [word2id.get(word, 1) if word != '\t' else word2id['<sep>'] for word in sentence]
        dec_input = torch.tensor(dec_input, dtype=torch.long, device=device).unsqueeze(0)

        output = self.greedy_decoder(dec_input).squeeze(0)
        out = [id2word[int(id)] for id in output]
        sep_indexs = []
        for i in range(len(out)):
            if out[i] == "<sep>":
                sep_indexs.append(i)
        answer = out[sep_indexs[-2] + 1:-1]
        answer = "".join(answer)
        return answer